﻿using System;
using System.Text;
using System.Security.Cryptography;

namespace System.Security.Cryptography
{
    //Generic PBKDF2 Class that can use any HMAC algorithm derived from the 
    // System.Security.Cryptography.HMAC abstract class

    // PER SPEC RFC2898 with help from user Dodgyrabbit on StackExchange
    // http://stackoverflow.com/questions/1046599/pbkdf2-implementation-in-c-sharp-with-rfc2898derivebytes

    // the use of default values for parameters in the functions puts this at .NET 4
    // if you remove those defaults and create the required constructors, you should be able to drop to .NET 2

    // USE AT YOUR OWN RISK!  I HAVE TESTED THIS AGAINST PUBLIC TEST VECTORS, BUT YOU SHOULD 
    // HAVE YOUR CODE PEER-REVIEWED AND SHOULD FOLLOW BEST PRACTICES WHEN USING CRYPTO-ANYTHING!
    // NO WARRANTY IMPLIED OR EXPRESSED, YOU ARE ON YOUR OWN!

    // PUBLIC DOMAIN!  NO COPYRIGHT INTENDED OR RESERVED!    

    //constrain T to be any class that derives from HMAC, and that exposes a new() constructor
    public class PBKDF2<T> : DeriveBytes where T : /* HMAC */ KeyedHashAlgorithm, new()
    {
        //Internal variables and public properties
        private int    _blockSize = -1;  // the byte width of the output of the HMAC algorithm      
        private T      _hmac;
        private byte[] _P = null;
        private int    _C = 0;
        private byte[] _S = null;

        // if you called the initializer/constructor specifying a salt size,
        // you will need this property to GET the salt after it was created from the crypto rng!
        // GET THIS BEFORE CALLING GETBYTES()!  OBJECT WILL BE RESET AFTER GETBYTES() AND
        // SALT WILL BE LOST!!
        public byte[] Salt { get { return (byte[])_S.Clone(); } }

        // Constructors
        public PBKDF2 (string Password, byte[] Salt, int IterationCount = 1000)
        { 
            Initialize(Password, Salt, IterationCount); 
        }

        public PBKDF2 (byte[] Password, byte[] Salt, int IterationCount = 1000)
        { 
            Initialize(Password, Salt, IterationCount); 
        }

        public PBKDF2 (string Password, int SizeOfSaltInBytes, int IterationCount = 1000)
        { 
            Initialize(Password, SizeOfSaltInBytes, IterationCount); 
        }

        public PBKDF2 (byte[] Password, int SizeOfSaltInBytes, int IterationCount = 1000)
        { 
            Initialize(Password, SizeOfSaltInBytes, IterationCount); 
        }

        //All Construtors call the corresponding Initialize methods
        private void Initialize (string Password, byte[] Salt, int IterationCount = 1000)
        {
            if (string.IsNullOrEmpty(Password))
                throw new ArgumentException("Password must contain meaningful characters and not be null.", "Password");

            if (IterationCount < 1)
                throw new ArgumentOutOfRangeException("IterationCount");

            Initialize(new UTF8Encoding(false).GetBytes(Password), Salt, IterationCount);
        }

        private void Initialize (byte[] Password, byte[] Salt, int IterationCount = 1000)
        {
            //all Constructors/Initializers eventually lead to this one which does all the "important" work
            if (Password == null || Password.Length == 0)
                throw new ArgumentException("Password cannot be null or empty.", "Password");

            if (Salt == null)
                Salt = new byte[0];

            if (IterationCount < 1)
                throw new ArgumentOutOfRangeException("IterationCount");

            _P = (byte[])Password.Clone();
            _S = (byte[])Salt.Clone();
            _C = IterationCount;
            
            //determine _blockSize
            
            _hmac       = new T();
            _hmac.Key   = new byte[] { 0 };
            byte[] test = _hmac.ComputeHash(new byte[] { 0 });
            _blockSize  = test.Length;

            Array.Clear(test, 0, test.Length);
        }

        private void Initialize (string Password, int SizeOfSaltInBytes, int IterationCount = 1000)
        {
            if (string.IsNullOrEmpty(Password))
                throw new ArgumentException("Password must contain meaningful characters and not be null.", "Password");

            if (IterationCount < 1)
                throw new ArgumentOutOfRangeException("IterationCount");

            Initialize(new UTF8Encoding(false).GetBytes(Password), SizeOfSaltInBytes, IterationCount);
        }

        private void Initialize (byte[] Password, int SizeOfSaltInBytes, int IterationCount = 1000)
        {
            if (Password == null || Password.Length == 0)
                throw new ArgumentException("Password cannot be null or empty.", "Password");

            if (SizeOfSaltInBytes < 0)
                throw new ArgumentOutOfRangeException("SizeOfSaltInBytes");

            if (IterationCount < 1)
                throw new ArgumentOutOfRangeException("IterationCount");

            // You didn't specify a salt, so I'm going to create one for you of the specific byte length
            byte[]                   data = new byte[SizeOfSaltInBytes];
            RNGCryptoServiceProvider rng  = new RNGCryptoServiceProvider();
            rng.GetBytes(data);
            // and then finish initializing...
            // Get the salt from the Salt parameter BEFORE calling GetBytes()!!!!!!!!!!!
            Initialize(Password, data, IterationCount);
        }

        /*
        ~PBKDF2()
        {
            //*DOOT* clean up in aisle 5! *KEKERKCRACKLE*
            this.Reset();
        }
        */

        // required by the Derive Bytes class/interface
        // this is where you request your output bytes after Initialize
        // state of class Reset after use!

        public override byte[] GetBytes (int ByteCount)
        {
            if (_S == null || _P == null)
                throw new InvalidOperationException("Object not Initialized!");

            if (ByteCount < 1)// || ByteCount > uint.MaxValue * blockSize)
                throw new ArgumentOutOfRangeException("ByteCount");

            int totalBlocks  = (int)Math.Ceiling((decimal)ByteCount / _blockSize);
            int partialBlock = (int)(ByteCount % _blockSize);
            byte[] result    = new byte[ByteCount];
            byte[] buffer    = null;

            // I'm using TT here instead of T from the spec because I don't want to confuse it with
            // the generic object T
            for (int TT = 1; TT <= totalBlocks; TT++)
            {
                // run the F function with the _C number of iterations for block number TT
                buffer = _F((uint)TT);
                //IF we're not at the last block requested
                //OR the last block requested is whole (not partial)
                //  then take everything from the result of F for this block number TT
                //ELSE only take the needed bytes from F
                if (TT != totalBlocks || (TT == totalBlocks && partialBlock == 0))
                    Buffer.BlockCopy(buffer, 0, result, _blockSize * (TT - 1), _blockSize);
                
                else
                    Buffer.BlockCopy(buffer, 0, result, _blockSize * (TT - 1), partialBlock);
            }

            //this.Reset();  // force cleanup after every use!  Cannot be reused!
            return result;
        }

        // required by the Derive Bytes class/interface
        public override void Reset()
        {
            _C = 0;

            if (_P != null)
            {
                _P.Initialize(); // the compiler might optimize this line out! :(
                _P = null;
            }

            if (_S != null)
            {
                _S.Initialize(); // the compiler might optimize this line out! :(
                _S = null;
            }

            if (_hmac != null)
                _hmac.Clear();

            _blockSize = -1;
        }

        // the core function of the PBKDF which does all the iterations
        // per the spec section 5.2 step 3
        private byte[] _F(uint I)
        {
            //NOTE: SPEC IS MISLEADING!!!
            //THE HMAC FUNCTIONS ARE KEYED BY THE PASSWORD! NEVER THE SALT!
            byte[] bufferU   = null;
            byte[] bufferOut = null;
            byte[] _int      = PBKDF2<T>.IntToBytes(I);

            _hmac     = new T();
            _hmac.Key = (_P); // KEY BY THE PASSWORD!

            _hmac.TransformBlock(_S, 0, _S.Length, _S, 0);
            _hmac.TransformFinalBlock(_int, 0, _int.Length);
            
            bufferU   = _hmac.Hash;
            bufferOut = (byte[])bufferU.Clone();

            for (int c = 1; c < _C; c++)
            {
                _hmac.Initialize();
                _hmac.Key = _P;  // KEY BY THE PASSWORD!
                bufferU   = _hmac.ComputeHash(bufferU);

                _Xor(ref bufferOut, bufferU);
            }

            return bufferOut;
        }

        // XOR one array of bytes into another (which is passed by reference)
        // this is the equiv of data ^= newData;
        private void _Xor (ref byte[] data, byte[] newData)
        {
            for (int i = data.GetLowerBound(0); i <= data.GetUpperBound(0); i++)
                data[i] ^= newData[i];
        }

        // convert an unsigned int into an array of bytes BIG ENDIEN
        // per the spec section 5.2 step 3
        static internal byte[] IntToBytes (uint i)
        {
            byte[] bytes = BitConverter.GetBytes(i);

            if (!BitConverter.IsLittleEndian)
                return bytes;

            else
            {
                Array.Reverse(bytes);
                return bytes;
            }
        }
    }
}
